import base64
import json
import os
import tempfile
from http import HTTPStatus
from pathlib import PurePosixPath
from urllib.parse import urlparse, unquote

import imagesize
import requests
from dashscope import ImageSynthesis
from fastapi import HTTPException

import config
from common.comm_utils import get_file_name
from server.openai_server import create_server

server_instance = create_server()

task = "image2image"
prompt_template = """
                你是一位精通DALL·E 3、Midjourney和Stable Diffusion的专业AI画图专家。你具有以下专长：
                1. 专业技能：
                - 深入理解各大AI绘画模型的特点和优势
                - 精通提示词工程，能够生成高质量、细节丰富的提示词
                - 熟悉各类艺术风格、构图技巧和视觉元素
                - 你需要根据以下公式生产提示词：
                提示词 = 主体（主体描述）+ 场景（场景描述）+ 风格（定义风格）+ 镜头语言 + 氛围词 + 细节修饰
                2. 工作方式：
                - 根据用户需求，生成结构化的提示词
                - 提供中英双语提示词（因为英文提示词通常效果更好）
                - 为不同平台优化提示词格式
                - 提供构图建议和参数设置
                3. 输出格式：
                [主题描述]：简单描述想要创建的图像
                [详细提示词]：
                - {{主体描述}}：确定主体清晰地描述图像中的主体，包括其特征、动作等。例如，"一个可爱的10岁中国小女孩，穿着红色衣服"。
                - {{场景描述}}：场景描述是对主体所处环境特征细节的描述，可通过形容词或短句列举。
                - {{定义风格}}：定义风格是明确地描述图像所应具有的特定艺术风格、表现手法或视觉特征。例如，"水彩风格"、"漫画风格"。
                - {{镜头语言}}：镜头语言包含景别、视角、渲染方式等。
                - {{光线氛围}}：光线氛围是对预期画面光线和氛围的描述，例如"梦幻"、"孤独"、"宏伟"，常见氛围词详见提示词词典。
                - {{细节修饰}}：细节修饰是对画面进一步的精细化和优化，以增强图像的细节表现力、丰富度和美感。例如"光源的位置"、"道具搭配"、"环境细节"，"高分辨率"等。
                [优化建议]：针对不同平台的特定优化建议
                请用专业且友好的语气与用户交流，理解他们的需求，并提供最适合的提示词建议。
                以下是用户提供的描述：
                {description}
                请根据用户的描述填充输出格式，不允许出现其他内容
                """


def description2params(description: str):
    prompt_template = """
                你是一位精通DALL·E 3、Midjourney和Stable Diffusion的专业AI画图专家：
                你需要输出n、size、seed、watermark、style、sketch_weight、sketch_extraction
                1.n integer [1-4] 必要
                生成图片的数量。取值范围为1~4张，默认为1。
                2.size string [768, 1440] 必要
                输出图像的分辨率。默认值是1024*1024。
                图像宽高边长的像素范围为：[768, 1440]，单位像素。
                以下是一些常用的图片分辨率：
                1024*768（XGA 4:3）
                1280*1024（5:4）
                1280*800（16:10）
                可任意组合以设置不同的图像分辨率，最高可达200万像素。
                3.seed的生成依据如下：int [0-2147483647] 可选
                seed代表随机数种子，用于控制模型生成内容的随机性。seed参数取值范围是(0, 2147483647)。
                如果不提供，则算法自动生成一个随机数作为种子。如果给定了，则根据n的值分别为n张图片生成seed参数。
                如果用户希望生成内容保持相对稳定，请使用相同的seed参数值。
                4.watermark的生成依据如下：true/false 可选
                是否添加水印标识，水印位于图片右下角，文案为"AI生成"。
                默认值为false，即不添加水印。
                5.style的格式如下： string 可选
                <auto>：默认值，由模型随机输出图像风格。
                <3d cartoon>：3D卡通。
                <anime>：二次元。
                <oil painting>：油画。
                <watercolor>：水彩。
                <sketch>：素描。
                <chinese painting>：中国画。
                <flat illustration>：扁平插画。
                如果用户明确了绘画风格，请根据用户的描述选择合适的style标签，否则请根据用户提供的描述信息选择最合适的style标签
                6.sketch_weight的生成依据如下：int [0-10] 可选
                输入草图对输出图像的约束程度。取值范围为0-10，取值间隔为1， 默认值为10。取值越大表示输出图像跟输入草图越相似。
                根据用户的描述判断其需要的相似度。
                7.sketch_extraction的生成依据如下：true/false 可选
                Boolean type, default value is False, indicating not to extract. Set to True when extracting.
                Determine whether the user needs to extract the sketch edge from the input image.
                以下是用户提供的描述：
                {description}
                输出示例：
                {{"n":1,"size":"768*768","style":"<auto>","seed":42,"sketch_weight":8,"sketch_extraction":false,"watermark":false}}
                不允许出现其他内容
                """
    prompt = prompt_template.format(description=description)
    response = server_instance.llm.invoke(prompt).content
    response_dict = json.loads(response)

    return response_dict


def get_image_size_imagesize(image_path):
    """获取图片尺寸
    
    Args:
        image_path: 可以是本地文件路径或URL
        
    Returns:
        tuple: (width, height)
    """
    try:
        if image_path.startswith(('http://', 'https://')):
            # 从URL中提取文件名
            file_name = PurePosixPath(unquote(urlparse(image_path).path)).parts[-1]
            # 检查本地是否存在该图片
            local_path = f'lucky_web/public/ai_material/ai_images/{file_name}'
            
            if os.path.exists(local_path):
                # 如果本地存在，直接获取本地文件尺寸
                width, height = imagesize.get(local_path)
                return width, height
            
            # 如果本地不存在，继续原有的URL处理逻辑
            response = requests.get(image_path)
            response.raise_for_status()
            
            with tempfile.NamedTemporaryFile(delete=False) as temp_file:
                temp_file.write(response.content)
                temp_path = temp_file.name
            
            try:
                # 获取临时文件的尺寸
                width, height = imagesize.get(temp_path)
                return width, height
            finally:
                # 清理临时文件
                try:
                    os.unlink(temp_path)
                except Exception as e:
                    print(f"清理临时文件失败: {str(e)}")
        else:
            # 如果是本地文件路径，直接获取尺寸
            width, height = imagesize.get(image_path)
            return width, height
            
    except Exception as e:
        print(f"获取图片尺寸失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"获取图片尺寸失败: {str(e)}")


def create_image(description: str, negative_prompt: str = None, sketch_image_url: str = None,extra_input:dict = None, mask_color=None, prompt_extend: bool = True):
    model=config.IMAGE_MODULE_NAME
    rsp = None
    params = description2params(description)
    n = params.get("n")
    size = params.get("size")
    seed = params.get("seed")
    watermark = params.get("watermark")
    if sketch_image_url:
        model = config.PROMPT2IMAGE_MODULE_NAME
        if mask_color is None:
            mask_color = []
        width, height = get_image_size_imagesize(sketch_image_url)
        style = params.get("style")
        sketch_weight = params.get("sketch_weight")
        sketch_extraction = params.get("sketch_extraction")
        rsp = ImageSynthesis.call(api_key=config.OPENAI_API_KEY,
                                  model=model,
                                  prompt=description,
                                  negative_prompt=negative_prompt,
                                  n=n,
                                  style=style,
                                  size=size,
                                  prompt_extend=prompt_extend,
                                  watermark=watermark,
                                  sketch_image_url=sketch_image_url,
                                  task=task,
                                  mask_color=mask_color,
                                  seed=seed,
                                  sketch_weight=sketch_weight,
                                  sketch_extraction=sketch_extraction,
                                  )
    elif extra_input:
        model = config.IMAGE2_NEW_IMAGE_MODULE_NAME
        rsp = ImageSynthesis.call(api_key=config.OPENAI_API_KEY,
                                  model=model,
                                  prompt=description,
                                  negative_prompt=negative_prompt,
                                  n=n,
                                  size=size,
                                  prompt_extend=prompt_extend,
                                  watermark=watermark,
                                  task=task,
                                  extra_input=extra_input,
                                  seed=seed,
                                  )
    else:
        rsp = ImageSynthesis.call(api_key=config.OPENAI_API_KEY,
                                  model=model,
                                  prompt=description,
                                  negative_prompt=negative_prompt,
                                  n=n,
                                  size=size,
                                  prompt_extend=prompt_extend,
                                  watermark=watermark,
                                  seed=seed,
                                  )
    print('response: %s' % rsp)
    if rsp.status_code == HTTPStatus.OK:
        if rsp.output.task_status == "FAILED":
            raise HTTPException(status_code=500, detail=rsp)
        try:
            # 在当前目录下保存图片
            for result in rsp.output.results:
                # file_name = PurePosixPath(unquote(urlparse(result.url).path)).parts[-1]
                file_name = get_file_name(description).lower()
                # 确保文件名有.wav后缀
                if not file_name.endswith('.png') and not file_name.endswith('.jpg') and not file_name.endswith('.jpeg') and not file_name.endswith('.webp'):
                    file_name += '.png'

                with open('lucky_web/public/ai_material/ai_images/%s' % file_name, 'wb+') as f:
                    f.write(requests.get(result.url).content)
                result_data= {'file_name': file_name,
                              'content': base64.b64encode(requests.get(result.url).content).decode('utf-8')}
                # 将二进制图片数据转换为base64字符串
                return result_data
        except Exception as e:
            raise HTTPException(status_code=500, detail=str(e))
    else:
        print('sync_call Failed, status_code: %s, code: %s, message: %s' %
              (rsp.status_code, rsp.code, rsp.message))
